package com.tianji.aigc.memory.jdbc;

import cn.hutool.core.collection.CollStreamUtil;
import cn.hutool.core.convert.Convert;
import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.tianji.aigc.entity.ChatRecord;
import com.tianji.aigc.memory.MessageUtil;
import com.tianji.aigc.service.IChatRecordService;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Repository;

import java.util.List;


/**
 * 基于MySQL实现的ChatMemoryRepository
 */
public class JdbcChatMemoryRepository implements ChatMemoryRepository {

    @Autowired
    private IChatRecordService iChatRecordService;

    @Override
    public List<String> findConversationIds() {
        /*List<ChatRecord> chatRecordList = this.iChatRecordService.lambdaQuery()
                .select(ChatRecord::getConversationId)
                .list();
        return CollStreamUtil.toList(chatRecordList,ChatRecord::getConversationId);
*/
        var queryWrapper = new QueryWrapper<ChatRecord>();
        queryWrapper.select("DISTINCT conversationId");
        var chatRecordList = this.iChatRecordService.list(queryWrapper);
        return CollStreamUtil.toList(chatRecordList, ChatRecord::getConversationId);
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        List<ChatRecord> chatRecordList = this.iChatRecordService.lambdaQuery()
                .eq(ChatRecord::getConversationId, conversationId)
                .orderByAsc(ChatRecord::getCreateTime)
                .list();
        return CollStreamUtil.toList(chatRecordList,chatRecord -> MessageUtil.toMessage(chatRecord.getData()));
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        Assert.notEmpty(messages, "消息列表不能为空");

        // 先删除该会话的所有历史消息
        this.deleteByConversationId(conversationId);

        // 通过对话id获取用户id
        var userId = Convert.toLong(StrUtil.subBefore(conversationId,'_',true));

        List<ChatRecord>  chatRecordList = CollStreamUtil.toList(messages,message -> ChatRecord.builder()
                .data(MessageUtil.toJson(message))
                .conversationId(conversationId)
                .creater(userId)
                .updater(userId)
                .build());
        this.iChatRecordService.saveBatch(chatRecordList);
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        // 删除会话消息
        var queryWrapper = Wrappers.<ChatRecord>lambdaQuery()
                .eq(ChatRecord::getConversationId,conversationId);
        this.iChatRecordService.remove(queryWrapper);
    }


}